package lab2;

import java.lang.reflect.InvocationTargetException;
import java.util.Comparator;
import java.util.Iterator;

import lab2.MyMap.Entry;

public class MyTreeMap<K, V> implements MyMap<K, V> {

	static class SimpleTreeEntry<K, V> implements MyMap.Entry<K, V> {

		K key;
		V value;
		
		SimpleTreeEntry<K, V> parent;
		SimpleTreeEntry<K, V> left;
		SimpleTreeEntry<K, V> right;
		
		boolean isRed;

        SimpleTreeEntry() {
        	this.key = null;
        	this.value = null;
        	parent = null;
        	left = null;
        	right = null;
        	isRed = false;
        }

        SimpleTreeEntry(final K key, final V value) {
        	this.key = key;
        	this.value = value;
        	parent = null;
        	left = null;
        	right = null;
        	isRed = false;
        }

        SimpleTreeEntry(final K key, final V value, final boolean isRed) {
        	this.key = key;
        	this.value = value;
        	parent = null;
        	left = null;
        	right = null;
        	this.isRed = isRed;
        }
        
        SimpleTreeEntry(final K key, final V value, final boolean isRed,
        		final Entry<K, V> parent, final Entry<K, V> left, final Entry<K, V> right) {
        	this.key = key;
        	this.value = value;
        	this.parent = (SimpleTreeEntry<K, V>) parent;
        	this.left = (SimpleTreeEntry<K, V>) left;
        	this.right = (SimpleTreeEntry<K, V>) right;
        	this.isRed = isRed;
        }

        public boolean equals(final Object o) {
			if (!(o instanceof SimpleTreeEntry<?, ?>)) {
				return false;
			}

			return (key == null ? ((SimpleTreeEntry<?, ?>) o).getKey() == null
					: key.equals(((SimpleTreeEntry<?, ?>) o).getKey())) &&
				     (value == null ? ((SimpleTreeEntry<?, ?>) o).getValue() == null
				     : value.equals(((SimpleTreeEntry<?, ?>) o).getValue()));
		}

		public int hashCode() {
			return (key == null ? 0 : key.hashCode()) ^
					(value == null ? 0 : value.hashCode());
		}

		public K getKey() {
			K res = key;
			return res;
		}

		public V getValue() {
			V res = value;
			return res;
		}

		public SimpleTreeEntry<K, V> getNext() {
			return left;
		}
		
		public SimpleTreeEntry<K, V> getLeft() {
			return left;
		}
		
		public SimpleTreeEntry<K, V> getRight() {
			return right;
		}

		public void setNext(final Entry<K, V> e) {
			left = (SimpleTreeEntry<K, V>) e;
		}
		
		public void setLeft(final Entry<K, V> e) {
			left = (SimpleTreeEntry<K, V>) e;
		}
		
		public void setRight(final Entry<K, V> e) {
			right = (SimpleTreeEntry<K, V>) e;
		}

		public int getHash() {
			return hashCode();
		}

		public V setValue(final V value) {
			V res = null;

			if (this.value != null) {
				res = this.value;
			}

			this.value = value;

			return res;
		}
		
		public void setIsRed(final boolean isRed) {
			this.isRed = isRed;
		}
		
		public String toString(){
			String res = "";
			
			if (isRed) {
				res += "R - ";
			} else {
				res += "B - ";
			}
			res += "Key: " + key.toString() + " | Value: " + value.toString();
			
			return res;
		}

	}
	
	final Comparator<? super K> myComparator;
	
	Entry<K, V> root;
	private int size;
	
	transient Entry<K, V>[] rootArray;
	private int rootArrayIndex = 0;
	
	public MyTreeMap() {
		myComparator = null;
		root = null;
		size = 0;
	}
	
	public MyTreeMap(final Comparator<? super K> comparator) {
		myComparator = comparator;
		root = null;
		size = 0;
	}
	
	
	public void clear() {
		root = null;
		size = 0;
	}

	public boolean containsKey(K key) {
		if (key == null) {
			throw new NullPointerException();
		}
		
		if (size == 0) {
			return false;
		}
		
		SimpleTreeEntry<K, V> p = (SimpleTreeEntry<K, V>) root;
		
		while (p != null) {
			if (myComparator == null) {
				if (!(p.getKey() instanceof Comparable)) {
					throw new ClassCastException("Keys can't be compared!");
				}
				Integer comparationResult = 0;

				try {
					comparationResult = (Integer) p.key.getClass().
							getMethod("compareTo", p.key.getClass()).
							invoke(p.key, key);
				} catch (IllegalAccessException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (IllegalArgumentException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (InvocationTargetException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (NoSuchMethodException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (SecurityException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
				
				if (comparationResult == 0) {
					return true;
				}
				
				if (comparationResult > 0) {
					p = p.left;
				} else {
					p = p.right;
				}
			} else {
				int comparationResult = myComparator.compare(p.key, key);
				
				if (comparationResult == 0) {
					return true;
				}
				
				if (comparationResult > 0) {
					p = p.left;
				} else {
					p = p.right;
				}
			}
		}
		return false;
	}

	public boolean containsValue(V value) {
		SimpleTreeEntry<K, V> p = (SimpleTreeEntry<K, V>) root;
		
		return searchForValue(p, value);		
	}	

	private boolean searchForValue(SimpleTreeEntry<K, V> p, V value) {
		
		if (size == 0) {
			return false;
		}
		
		if (p == null) {
			return false;
		}
						
		if (!(p.getValue() instanceof Comparable)) {
			throw new ClassCastException("Values can't be compared!");
		}
		Integer comparationResult = 0;

		try {
			comparationResult = (Integer) p.value.getClass().
					getMethod("compareTo", p.value.getClass()).
					invoke(p.value, value);
		} catch (IllegalAccessException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IllegalArgumentException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InvocationTargetException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchMethodException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (SecurityException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

		if (comparationResult == 0) {
			return true;
		}

		return (searchForValue(p.left, value) || searchForValue(p.right, value));
	}

	public V get(K key) {
		SimpleTreeEntry<K, V> p = (SimpleTreeEntry<K, V>) root;
		
		return getValueForKey(p, key);
	}

	private V getValueForKey(SimpleTreeEntry<K, V> p, K key) {
		
		if (size == 0) {
			return null;
		}
		
		if (p == null) {
			return null;
		}
		
		while (p != null) {
			if (myComparator == null) {
				if (!(p.key instanceof Comparable)) {
					throw new ClassCastException("Keys can't be compared!");
				}
				Integer comparationResult = 0;

				try {
					comparationResult = (Integer) p.key.getClass().
							getMethod("compareTo", p.key.getClass()).
							invoke(p.key, key);
				} catch (IllegalAccessException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (IllegalArgumentException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (InvocationTargetException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (NoSuchMethodException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (SecurityException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
				
				if (comparationResult == 0) {
					return p.value;
				}
				
				if (comparationResult > 0) {
					p = p.left;
				} else {
					p = p.right;
				}
			} else {
				int comparationResult = myComparator.compare(p.key, key);
				
				if (comparationResult == 0) {
					return p.value;
				}
				
				if (comparationResult > 0) {
					p = p.left;
				} else {
					p = p.right;
				}
			}
		}
		return null;
	}
	
	private SimpleTreeEntry<K, V> getPlaceForKey(SimpleTreeEntry<K, V> p, K key) {
		
		if (p == null) {
			return p;
		}
		
		while (p != null) {
			if (myComparator == null) {
				if (!(p.key instanceof Comparable)) {
					throw new ClassCastException("Keys can't be compared!");
				}
				Integer comparationResult = 0;

				try {
					comparationResult = (Integer) p.key.getClass().
							getMethod("compareTo", p.key.getClass()).
							invoke(p.key, key);
				} catch (IllegalAccessException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (IllegalArgumentException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (InvocationTargetException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (NoSuchMethodException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (SecurityException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
				
				if (comparationResult == 0) {
					return p;
				}
				
				if (comparationResult > 0) {
					p = p.left;
				} else {
					p = p.right;
				}
			} else {
				int comparationResult = myComparator.compare(p.key, key);
				
				if (comparationResult == 0) {
					return p;
				}
				
				if (comparationResult > 0) {
					p = p.left;
				} else {
					p = p.right;
				}
			}
		}
		return p;
	}

	private SimpleTreeEntry<K, V> getParentForKey(SimpleTreeEntry<K, V> p, K key) {
	
		if (p == null) {
			return p;
		}

		while (p != null) {
			if (myComparator == null) {
				if (!(p.key instanceof Comparable)) {
					throw new ClassCastException("Keys can't be compared!");
				}
				Integer comparationResult = 0;

				try {
					comparationResult = (Integer) p.key.getClass().
							getMethod("compareTo", p.key.getClass()).
							invoke(p.key, key);
				} catch (IllegalAccessException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (IllegalArgumentException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (InvocationTargetException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (NoSuchMethodException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (SecurityException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}

				if (comparationResult == 0) {
					return p;
				}

				if (comparationResult > 0) {
					if (p.left == null) {
						return p;
					}
					p = p.left;
				} else {
					if (p.right == null) {
						return p;
					}
					p = p.right;
				}
			} else {
				int comparationResult = myComparator.compare(p.key, key);

				if (comparationResult == 0) {
					return p;
				}

				if (comparationResult > 0) {
					if (p.left == null) {
						return p;
					}
					p = p.left;
				} else {
					if (p.right == null) {
						return p;
					}
					p = p.right;
				}
			}
		}
		return null;
	}

	public boolean isEmpty() {
		return (size == 0);
	}

	public V put(K key, V value) {
		if (isEmpty()) {
			root = new SimpleTreeEntry<K, V>(key, value, false);
			size++;
			return null;
		}
		
		SimpleTreeEntry<K, V> p = getPlaceForKey((SimpleTreeEntry<K, V>) root, key);
		
		if (p != null) {
			V res = p.value;
			p.value = value;
			return res;
		}
		
		SimpleTreeEntry<K, V> q = getParentForKey((SimpleTreeEntry<K, V>) root, key);
		
		p = new SimpleTreeEntry<K, V>(key, value, true, q, null, null);
		
		if (myComparator == null) {
			if (!(p.key instanceof Comparable)) {
				throw new ClassCastException("Keys can't be compared!");
			}
			Integer comparationResult = 0;

			try {
				comparationResult = (Integer) p.key.getClass().
						getMethod("compareTo", p.key.getClass()).
						invoke(p.key, q.key);
			} catch (IllegalAccessException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (IllegalArgumentException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (InvocationTargetException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (NoSuchMethodException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (SecurityException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
			
			if (comparationResult < 0) {
				q.left = p;
			} else {
				q.right = p;
			}
		} else {
			int comparationResult = myComparator.compare(p.key, q.key);
			
			if (comparationResult < 0) {
				q.left = p;
			} else {
				q.right = p;
			}
		}
		
		size++;
		
		arrangeTreeColors(p);
		
		return null;
	}
	
	private SimpleTreeEntry<K, V> getGrandParentForEntry(SimpleTreeEntry<K, V> p) {
		if ((p != null) && (p.parent != null)) {
			return p.parent.parent;
		} else {
			return null;
		}
	}
	
	private SimpleTreeEntry<K, V> getUncleForEntry(SimpleTreeEntry<K, V> p) {
		SimpleTreeEntry<K, V> g = getGrandParentForEntry(p);
		if (g == null) {
			return null;
		}
		if (p.parent == g.left) {
			return g.right;
		} else {
			return g.left;
		}
	}

	private void arrangeTreeColors(SimpleTreeEntry<K, V> p) {
		if (p.parent == null) {
			p.isRed = false;
			return;
		}
		
		if (!p.parent.isRed) {
			return;
		}
		
		SimpleTreeEntry<K, V> u = getUncleForEntry(p);
		SimpleTreeEntry<K, V> g = getGrandParentForEntry(p);
		
		if ((u != null) && (u.isRed)) {
			  p.parent.isRed = false;
			  u.isRed = false;
			  g = getGrandParentForEntry(p);
			  g.isRed = true;
			  arrangeTreeColors(g);
			  return;
		}
		
		if ((p == p.parent.right) && (p.parent == g.left)) {
			  rotateLeft((SimpleTreeEntry<K, V>) root, p.parent);
			  p = p.left;
		} else if ((p == p.parent.left) && (p.parent == g.right)) {
			  rotateRight((SimpleTreeEntry<K, V>) root, p.parent);
			  p = p.right;
		}
		
		g = getGrandParentForEntry(p);
		p.parent.isRed = false;
		g.isRed = true;
		if (p == p.parent.left)
			rotateRight((SimpleTreeEntry<K, V>) root, g);
		else
			rotateLeft((SimpleTreeEntry<K, V>) root, g);
	}

	private void rotateLeft(SimpleTreeEntry<K, V> root, SimpleTreeEntry<K, V> x) {
		SimpleTreeEntry<K, V> y = x.right;
		x.right = y.left;
		if (y.left != null) {
			y.left.parent = x;
		}
		
		y.parent = x.parent;
		
		if (x.parent == null) {
			root = y;
		} else {
			if (x == x.parent.left) {
				x.parent.left = y;
			} else {
				x.parent.right = y;
			}
		}
		
		y.left = x;
		x.parent = y;
	}
	
	private void rotateRight(SimpleTreeEntry<K, V> root, SimpleTreeEntry<K, V> y) {
		SimpleTreeEntry<K, V> x = y.left;
		y.left = x.right;
		if (x.right != null) {
			x.right.parent = y;
		}
		
		x.parent = y.parent;
		
		if (y.parent == null) {
			root = x;
		} else {
			if (y == y.parent.left) {
				y.parent.left = x;
			} else {
				y.parent.right = x;
			}
		}
		
		x.right = y;
		y.parent = x;
	}

	public V remove(K key) {
		if (isEmpty()) {
			return null;
		}
		
		SimpleTreeEntry<K, V> p = getPlaceForKey((SimpleTreeEntry<K, V>) root, key);
		
		if (p == null) {
			return null;
		}
		
		V res = p.value;
		removeEntry(p);
		size--;
		return res;
	}

	private void removeEntry(SimpleTreeEntry<K, V> p) {
		boolean pWasRed = p.isRed;
		boolean qWasRed;
		
		if (p.left == null && p.right == null) {
			if (p.parent.left == p) {
				p.parent.left = null;
				return;
			} else {
				p.parent.right = null;
				return;
			}
		} else {
			if (p.left == null) {
				if (p.parent.left == p) {
					p.parent.left = p.right;
					return;
				} else {
					p.parent.right = p.right;
					return;
				}
			} else {
				if (p.right == null) {
					if (p.parent.left == p) {
						p.parent.left = p.left;
						return;
					} else {
						p.parent.right = p.left;
						return;
					}
				} else {
					SimpleTreeEntry<K, V> q = findImmediateSuccessor(p);
					p.key = q.key;
					p.value = q.value;
					p.isRed = q.isRed;
					qWasRed = q.isRed;
					removeEntry(q);
				}
			}
		}
		
		if (!pWasRed && qWasRed) {
			SimpleTreeEntry<K, V> r = (SimpleTreeEntry<K, V>) root;
			arrangeTreeColors(r);
			return;
		} else {
			doubleBlackElimination(p);
		}		
	}

	private SimpleTreeEntry<K, V> findImmediateSuccessor(SimpleTreeEntry<K, V> p) {
		SimpleTreeEntry<K, V> is = p.right;
		while (is.left != null) {
			is = is.left;
		}
		return is;
	}
	
	private void doubleBlackElimination(SimpleTreeEntry<K, V> p) {
		if (p.parent == null) {
			return;
		}
		
		SimpleTreeEntry<K, V> sib = findSibling(p);
		
		if (sib.isRed) {
			p.parent.isRed = true;
			sib.isRed = false;
			if (p == p.parent.left) {
				rotateLeft((SimpleTreeEntry<K, V>) root, p.parent);
			} else {
				rotateRight((SimpleTreeEntry<K, V>) root, p.parent);
			}			
		}
		
		if (!p.parent.isRed && !sib.isRed && !sib.left.isRed && !sib.right.isRed) {
			sib.isRed = true;
			doubleBlackElimination(p.parent);
			return;
		}
		
		if (p.parent.isRed && !sib.isRed && !sib.left.isRed && !sib.right.isRed) {
			sib.isRed = true;
			p.parent.isRed = false;
			return;
		}
		
		if (!sib.isRed) {
			if ((p == p.parent.left) && !sib.right.isRed && sib.left.isRed) {
				sib.isRed = true;
				sib.left.isRed = false;
				rotateRight((SimpleTreeEntry<K, V>) root, sib);
			} else {
				if ((p == p.parent.right) && sib.right.isRed && !sib.left.isRed) {
					sib.isRed = true;
					sib.right.isRed = false;
					rotateLeft((SimpleTreeEntry<K, V>) root, sib);
				}
			}
		}
		
		sib.isRed = p.parent.isRed;
		p.parent.isRed = false;
		if (p == p.parent.left) {
			sib.right.isRed = false;
			rotateLeft((SimpleTreeEntry<K, V>) root, p.parent);			
		} else {
			sib.left.isRed = false;
			rotateRight((SimpleTreeEntry<K, V>) root, p.parent);
		}
	}
	
	private SimpleTreeEntry<K, V> findSibling(SimpleTreeEntry<K, V> p) {
		if (p == p.parent.left) {
			return p.parent.right;
		} else {
			return p.parent.left;
		}
	}

	public int size() {
		return size;
	}

	public Iterator<lab2.MyMap.Entry<K, V>> entryIterator() {
		Iterator<Entry<K, V>> it = new Iterator<Entry<K, V>>() {
			
			Entry<K, V>[] p = toArray((SimpleTreeEntry<K, V>) root);
			int index = 0;
			
			public boolean hasNext() {
                return (index < size);
            }

            public Entry<K, V> next() {
            	Entry<K, V> res = p[index];
            	index++;
                return res;
            }

            public void remove() {
                // TODO Auto-generated method stub
            }
        };
        return it;
	}
	
	protected Entry<K, V>[] toArray(SimpleTreeEntry<K, V> root) {
		if (isEmpty()) {
			return null;
		}
		
		rootArray = new Entry[size];
		rootArrayIndex = 0;
		addToArray(root);

		return rootArray;
	}

	private void addToArray(SimpleTreeEntry<K, V> p) {
		rootArray[rootArrayIndex] = p;
		if (p.left != null) {
			rootArrayIndex++;
			addToArray(p.left);
		}
		if (p.right != null) {
			rootArrayIndex++;
			addToArray(p.right);
		}
	}

	public String toString() {				
		return entryToString((SimpleTreeEntry<K, V>) root);
	}
	
	private String entryToString(SimpleTreeEntry<K, V> e) {
		String res = "";

		res += e.toString();
		res += " ( ";
		if (e.left != null) {
			res += entryToString(e.left) + " ; ";
		}
		else {
			res += " --- ; ";
		}		
		if (e.right != null) {
			res += entryToString(e.right);
		}
		else {
			res += " --- ";
		}
		res += " ) ";

		return res;
	}
}
